package piggy.pool;

import io.netty.channel.Channel;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

public class DynamicConnectionPool {
    private final String addr;
    private final int minConnections;
    private final int maxConnections;
    private final long connectionTimeout;
    private final BlockingQueue<Channel> connectionQueue;
    private final Set<Channel> activeConnections;
    private final ScheduledExecutorService healthCheckScheduler;
    private final AtomicInteger currentConnections; // 当前连接数

    // 构造函数初始化连接池
    public DynamicConnectionPool(String addr, int minConnections, int maxConnections, long connectionTimeout) {
        this.addr = addr;
        this.minConnections = minConnections;
        this.maxConnections = maxConnections;
        this.connectionTimeout = connectionTimeout;

        // 队列用于存放空闲连接
        this.connectionQueue = new LinkedBlockingQueue<>(maxConnections);
        this.activeConnections = Collections.synchronizedSet(new HashSet<>());
        this.currentConnections = new AtomicInteger(0);

        // 初始化连接池
        initializeConnections();

        // 启动健康检查任务
        healthCheckScheduler = Executors.newScheduledThreadPool(1);
    }

    // 初始化连接池，创建最小连接数的连接
    private void initializeConnections() {
        for (int i = 0; i < minConnections; i++) {
            createAndAddConnection();
        }
    }

    // 创建连接并添加到连接池
    private void createAndAddConnection() {
        if (currentConnections.get() < maxConnections) {
            Channel channel = NettyClient.getChannel(addr);
            boolean offer = connectionQueue.offer(channel);
            currentConnections.incrementAndGet();
        }
    }

    // 获取一个数据库连接
    public Channel getConnection() throws InterruptedException {
        // 获取连接，超时机制
        Channel connection = connectionQueue.poll(connectionTimeout, TimeUnit.MILLISECONDS);
        if (connection == null) {
            // 如果连接池为空并且当前连接数小于最大连接数，则创建新连接
            if (currentConnections.get() < maxConnections) {
                createAndAddConnection();
                connection = connectionQueue.poll(connectionTimeout, TimeUnit.MILLISECONDS);
            }
        }

        // 如果连接仍然为 null，则抛出异常
        if (connection == null) {
            throw new RuntimeException("Unable to get connection from the pool within the timeout period.");
        }

        // 将连接标记为活动连接
        activeConnections.add(connection);
        return connection;
    }

    // 归还连接到连接池
    public void releaseConnection(Channel connection) {
        if (connection != null) {
            // 如果连接有效，则将其返回连接池
            if (!connection.isActive()) {
                activeConnections.remove(connection);
                if (connectionQueue.size() < maxConnections) {
                    boolean offer = connectionQueue.offer(connection);// 将连接归还到池中
                } else {
                    connection.close();  // 如果连接池满了，则关闭连接
                }
            } else {
                activeConnections.remove(connection);
                connection.close();
            }
        }
    }

    // 关闭连接池中的所有连接
    public void shutdown() {
        healthCheckScheduler.shutdown();
        for (Channel connection : connectionQueue) {
            connection.close();
        }
    }


    // 动态调整连接池大小
    public void adjustPoolSize() {
        int activeCount = activeConnections.size();
        int idleCount = connectionQueue.size();
        // 如果空闲连接数过多，回收连接
        if (idleCount > minConnections) {
            int excess = idleCount - minConnections;
            for (int i = 0; i < excess; i++) {
                Channel connection = connectionQueue.poll();
                if (connection != null) {
                    connection.close();
                    currentConnections.decrementAndGet();
                }
            }
        }
        // 如果活动连接数接近最大连接数，增加连接
        if (activeCount < minConnections && currentConnections.get() < maxConnections) {
            createAndAddConnection();
        }
    }
}

